# -*- coding: utf-8 -*-
import os
import sys
import warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# 将自定义 pymagic 路径加入环境变量
pymagic_path = '/database/home/duansizhang/pywave'
if pymagic_path not in sys.path:
    sys.path.insert(0, pymagic_path)

from pymagic.tools.ecgdllinterface import DataSource, EcgAnalyzer
from hrvanalysis import (
    remove_outliers, interpolate_nan_values, remove_ectopic_beats,
    get_time_domain_features, get_frequency_domain_features,
    get_geometrical_features, get_sampen
)
import EntropyHub

warnings.filterwarnings('ignore')

# --- 配置 ---
raw_file  = '/database/private/mgcdb/raw/0a0b34b8220644848209d7a199afd6fc_1699056000000.raw'
out_dir   = '/database/home/duansizhang/hrv_predict/result'
csv_global= os.path.join(out_dir, 'hrv_global_window.csv')
csv_rr    = os.path.join(out_dir, 'hrv_rr_features.csv')
img_rr    = os.path.join(out_dir, 'rr_plot.png')
fs        = 250     # 采样率
win_sec   = 300     # 窗口长度（秒）

os.makedirs(out_dir, exist_ok=True)

# 1. 读取 ECG 与 R 峰
ds        = DataSource(raw_file, 3)
res       = EcgAnalyzer(ds)
rpos_all  = np.array(res['rpos'], dtype=int)
anntype   = np.array(res['anntype'], dtype=int)

# 只保留真实 R 波
rpos      = rpos_all[anntype == 1]
times     = rpos / fs

# 2. 计算 RR 间期(ms)
rr_ms     = np.diff(rpos) * 1000.0 / fs
rr_times  = times[1:]

# 3. 清洗
rr_clean  = remove_outliers(rr_intervals=rr_ms, low_rri=300, high_rri=2000)
rr_interp = interpolate_nan_values(rr_intervals=rr_clean)
nn_beats  = remove_ectopic_beats(rr_intervals=rr_interp)
nn_interp = interpolate_nan_values(rr_intervals=nn_beats)
mask      = ~np.isnan(nn_interp)
nn_array  = np.array(nn_interp)[mask]
nn_times  = rr_times[mask]

# 4. 特征提取函数
def extract_features(nn_seg):
    if len(nn_seg) < 3:
        return None
    feats = {}
    td = get_time_domain_features(nn_seg)
    fd = get_frequency_domain_features(nn_seg)
    geo= get_geometrical_features(nn_seg)
    feats.update({k: float(td[k])   if td[k]   is not None else np.nan for k in td})
    feats.update({k: float(fd[k])   if fd[k]   is not None else np.nan for k in fd})
    feats.update({k: float(geo[k])  if geo[k]  is not None else np.nan for k in geo})
    if len(nn_seg) > 10:
        smen = get_sampen(nn_seg)
        feats.update({k: float(smen[k]) if smen[k] is not None else np.nan for k in smen})
        try:
            apen_vals, _ = EntropyHub.ApEn(nn_seg)
            feats['ApEn'] = float(apen_vals[2]) if len(apen_vals) > 2 else np.nan
        except:
            feats['ApEn'] = np.nan
    else:
        feats['ApEn'] = np.nan
    feats['nbeats']   = len(nn_seg)
    feats['mean_hr']  = 60000.0 / np.mean(nn_seg) if np.mean(nn_seg) > 0 else np.nan
    return feats

# 5. 构建 rows
rows = []

# 5.1 全局
g_feats = extract_features(nn_array)
if g_feats:
    rows.append({
        'label': 'global',
        'window_start': 0.0,
        'window_end': float(nn_times[-1]),
        'timestamp': 0.0,
        **g_feats
    })

# 5.2 滑窗
t_end = float(nn_times[-1])
for ws in np.arange(0, t_end, win_sec):
    we = ws + win_sec
    idx = (nn_times >= ws) & (nn_times < we)
    seg = nn_array[idx]
    feats = extract_features(seg)
    if feats:
        rows.append({
            'label': 'window',
            'window_start': float(ws),
            'window_end': float(we),
            'timestamp': float(ws + win_sec / 2),
            **feats
        })

# 6. 保存 CSV
df = pd.DataFrame(rows)
df.to_csv(csv_global, index=False, encoding='utf-8-sig')
if g_feats:
    pd.DataFrame([g_feats]).to_csv(csv_rr, index=False, encoding='utf-8-sig')

# 7. 绘制 RR 间期图
plt.figure(figsize=(10,3))
plt.plot(nn_times, nn_array, '-o')
plt.xlabel('Time (s)')
plt.ylabel('RR (ms)')
plt.grid(True)
plt.tight_layout()
plt.savefig(img_rr, dpi=300)
plt.close()
